package us.luosl.lslt.function.stream;

import us.luosl.lslt.lang.Tuple2;

import java.util.*;
import java.util.function.*;
import java.util.stream.Stream;

public interface EntryStreamWrapper<K, V> extends StreamWrapper<Tuple2<K, V>> {

    default <T1, T2> EntryStreamWrapper<T1, T2> wrap(Stream<Tuple2<T1, T2>> stream) {
        return () -> stream;
    }

    default GeneralStreamWrapper<Tuple2<K, V>> toGeneral(){
        return StreamWrapper.wrap(getStream());
    }

    default EntryStreamWrapper<K, V> filter(BiPredicate<K, V> predicate) {
        return wrap(getStream().filter(tuple -> tuple.test(predicate)));
    }

    default <T1, T2> EntryStreamWrapper<T1, T2> map(BiFunction<K, V, Tuple2<T1, T2>> mapper) {
        return wrap(getStream().map(tuple -> tuple.map(mapper)));
    }

    default <R> GeneralStreamWrapper<R> mapToGeneral(BiFunction<K, V, R> mapper) {
        return StreamWrapper.wrap(getStream().map(t -> t.apply(mapper)));
    }

    default <R> EntryStreamWrapper<R, V> mapKeys(Function< K, R> mapper) {
        return wrap(getStream().map(tuple -> tuple.map1(mapper)));
    }

    default <R> EntryStreamWrapper<K, R> mapValues(Function<V, R> mapper) {
        return wrap(getStream().map(tuple -> tuple.map2(mapper)));
    }

    default GeneralStreamWrapper<K> keyStreamWrapper(){
        return mapToGeneral((k, v) -> k);
    }

    default GeneralStreamWrapper<V> valueStreamWrapper(){
        return mapToGeneral((k, v) -> v);
    }

    default <T1, T2> EntryStreamWrapper<T1, T2> flatMap(BiFunction<K, V, EntryStreamWrapper<T1, T2>> mapper) {
        return wrap(getStream().flatMap(t -> t.apply(mapper).getStream()));
    }

    default void forEach(BiConsumer<K, V> action) {
        getStream().forEach(tuple -> tuple.accept(action));
    }

    default void forEachOrdered(BiConsumer<K, V> action) {
        getStream().forEachOrdered(tuple -> tuple.accept(action));
    }

    default EntryStreamWrapper<K, V> distinct() {
        return wrap(getStream().distinct());
    }

    default EntryStreamWrapper<K, V> sorted() {
        return wrap(getStream().sorted());
    }

    default EntryStreamWrapper<K, V> sorted(Comparator<? super Tuple2<K, V>> comparator) {
        return wrap(getStream().sorted(comparator));
    }

    default EntryStreamWrapper<K, V> peek(BiConsumer<K, V> action) {
        return wrap(getStream().peek(tuple -> tuple.accept(action)));
    }

    default EntryStreamWrapper<K, V> limit(long maxSize) {
        return wrap(getStream().limit(maxSize));
    }

    default EntryStreamWrapper<K, V> skip(long n) {
        return wrap(getStream().skip(n));
    }

    default Map<K, V> toMap(){
        return toMap(Tuple2::get_1, Tuple2::get_2);
    }

    default EntryStreamWrapper<K, V> sortBy(Supplier<Comparator<Tuple2<K, V>>> supplier, boolean nullFirst){
        Comparator<Tuple2<K, V>> comparator = supplier.get();
        if(nullFirst){
            comparator = Comparator.nullsFirst(comparator);
        }else{
            comparator = Comparator.nullsLast(comparator);
        }
        return sorted(comparator);
    }

    default <R extends Comparable<R>> EntryStreamWrapper<K, V> sortBy(BiFunction<K, V, R> orderExtractor, boolean nullFirst) {
        Function<Tuple2<K, V>, R> extractor = tuple -> tuple.apply(orderExtractor);
        return sortBy(() -> Comparator.comparing(extractor), nullFirst);
    }

    default <R> EntryStreamWrapper<K, V> sortBy(BiFunction<K, V, R> orderExtractor, Comparator<R> byComparator, boolean nullFirst) {
        Function<Tuple2<K, V>, R> extractor = tuple -> tuple.apply(orderExtractor);
        return sortBy(() -> Comparator.comparing(extractor, byComparator), nullFirst);
    }

    default <R extends Comparable<R>> EntryStreamWrapper<K, V> sortByDesc(BiFunction<K, V, R> orderExtractor, boolean nullFirst) {
        Function<Tuple2<K, V>, R> extractor = tuple -> tuple.apply(orderExtractor);
        return sortBy(() -> Comparator.comparing(extractor).reversed(), nullFirst);
    }

    default <R> EntryStreamWrapper<K, V> sortByDesc(BiFunction<K, V, R> orderExtractor, Comparator<R> byComparator, boolean nullFirst) {
        Function<Tuple2<K, V>, R> extractor = tuple -> tuple.apply(orderExtractor);
        return sortBy(() -> Comparator.comparing(extractor, byComparator).reversed(), nullFirst);
    }

    default EntryStreamWrapper<K, V> reverse() {
        return StreamWrapper.wrap(getStream()).reverse().mapToEntry(i -> i);
    }

    default <U> U foldRight(U zero, BiFunction<U, Tuple2<K, V>, U> f) {
        return reverse().foldLeft(zero, f);
    }

    default Map<K, List<V>> groupByKey(){
        return agg(Tuple2::get_1, list -> StreamWrapper.wrapEntry(list).valueStreamWrapper().toList());
    }

    default <KEY> EntryStreamWrapper<K, V> distinctBy(BiFunction<K, V, KEY> byExtractor){
        final Set<KEY> distSet = new HashSet<>();
        return wrap(getStream().filter(item -> distSet.add(item.apply(byExtractor))));
    }

    default EntryStreamWrapper<K, V> filterNot(BiPredicate<K, V> predicate){
        return filter(predicate.negate());
    }

}
